import tensorflow as tf
from stock_a.feature import feature_pb2
from stock_a.feature.abstract_feature import AbstractFeature
from stock_a.feature.stock_feature_source import StockFeatureSource
from stock_a.feature.global_feature_source import GlobalFeatureSource
from stock_a.common.db import stock_info


class IdFeature(AbstractFeature):
    def __init__(self, stock_feature_source: StockFeatureSource, global_feature_source: GlobalFeatureSource):
        super(IdFeature, self).__init__(stock_feature_source, global_feature_source)

    def feature_name(self):
        return "id"

    def feature_type(self):
        return tf.int64

    def feature_shape(self):
        return ()

    def set_pb_feature(self, idx: int, feature_pb: feature_pb2.feature):
        code = self.stock_feature_source.day_k_info_list[idx][stock_info.CODE_INDEX]
        stock_id = self.global_feature_source.code_to_id_map[code]
        if stock_id is None:
            raise Exception('code={}, but id is none'.format(code))
        feature_pb.id = stock_id

    def get_pb_feature(self, feature_pb: feature_pb2.feature):
        return [feature_pb.id]